"""
python fujian1_autoArima_predict.py
"""

import os
import json
import pandas as pd
import pmdarima as pm
from datetime import datetime
from dateutil.relativedelta import relativedelta

# 文件路径
input_dir = os.path.join("fujian", "fujian1", "cubic_spline", "interpolation_output")
output_dir = os.path.join("fujian", "fujian1", "auto_arima", "result")
log_dir = os.path.join(output_dir, "log")

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)
os.makedirs(log_dir, exist_ok=True)

# 总日志文件路径
all_log_path = os.path.join(log_dir, "all_logs.log")

# 日期范围
start_date = datetime(2022, 7, 1)
forecast_dates = [start_date + relativedelta(months=i) for i in range(15)]

# 存储有问题的category
problematic_categories = []

# 遍历每个输入文件
# ... (previous code remains unchanged)

# 遍历每个输入文件
for file_name in os.listdir(input_dir):
    if file_name.endswith(".json") and "interpolation_categorycategory" in file_name:
        category_id = file_name.replace("interpolation_categorycategory", "").replace(".json", "")
        
        # 加载库存数据
        with open(os.path.join(input_dir, file_name), "r") as f:
            data = json.load(f)
        
        # 提取日期和库存量
        dates = [item["date"] for item in data]
        inventory = [item["inventory"] for item in data]
        
        # 将数据转换为 Pandas Series
        inventory_series = pd.Series(inventory, index=pd.to_datetime(dates)).asfreq('MS')

        # 创建日志文件
        log_file_path = os.path.join(log_dir, f"category_{category_id}.log")
        with open(log_file_path, "w") as log_file:
            # 检查 NaN 值
            if inventory_series.isnull().sum() > 0:
                log_file.write(f"Error: {inventory_series.isnull().sum()} NaN values found in category {category_id} data.\n")
                problematic_categories.append(f"Category {category_id} has NaN values.")
                continue  # Skip this category if NaN values exist

            log_file.write(f"Data for category {category_id} is normal, with no missing values.\n")
        
        try:
            # 使用 auto_arima 生成最佳参数并进行模型拟合
            model = pm.auto_arima(inventory_series, start_p=0, start_q=0, max_p=3, max_q=3,
                                stepwise=True, maxiter=300,
                                error_action='ignore', suppress_warnings=True)

            # Check for NaN values right before prediction
            if inventory_series.isnull().sum() > 0:
                raise ValueError("NaN values detected in inventory series before prediction.")

            # 预测未来三个月
            forecast = model.predict(n_periods=3)

            # 将预测数据加入原始数据
            for i, pred in enumerate(forecast, start=12):
                inventory_series[forecast_dates[i]] = pred

            # 生成输出格式
            output_data = [{"date": date.strftime("%Y-%m-%d"), "inventory": int(inventory_series[date])} for date in forecast_dates]

            # 保存到输出文件
            output_file = os.path.join(output_dir, f"autoArima_category{category_id}.json")
            with open(output_file, "w") as f:
                json.dump(output_data, f, indent=4)

            print(f"完成 {category_id} 类别的预测，结果已保存到 {output_file}")

        except ValueError as e:
            error_message = f"Error during prediction for category {category_id}: {str(e)}"
            log_file.write(error_message + "\n")
            problematic_categories.append(error_message)  # Add to the problematic categories list

# 输出所有问题分类到总日志文件
with open(all_log_path, "w") as all_log_file:
    if problematic_categories:
        all_log_file.write("Problematic Categories:\n")
        for category in problematic_categories:
            all_log_file.write(category + "\n")
    else:
        all_log_file.write("No problematic categories found.\n")

print(f"所有问题分类的汇总已保存到 {all_log_path}")
